import onnx
import onnxruntime
import numpy as np
import torch
from torchvision.transforms import transforms
import h5py
from PIL import Image
import time

# model = onnx.load("BiFNet_32_44186.onnx")
# onnx.checker.check_model(model)
#
# session = onnxruntime.InferenceSession('BiFNet_32_44186.onnx')
# input_names = [inp.name for inp in session.get_inputs()]
# print(input_names)
# class_images = np.random.randn(2, 1, 224, 224).astype(np.float32)
# defect_images = np.random.randn(2, 1, 224, 224).astype(np.float32)
# inputs_dict = {
#     'class_images': class_images,
#     'defect_images': defect_images
# }
# outputs = session.run(None, inputs_dict)
# print(outputs[0].shape)



def predictClassDefectImages(class_image, defect_image, onnx_weight):
    categories = ['ADI_particle_developed', 'Array_peeling', 'Cu_missing', 'Other_peeling', 'Partial_etch',
                  'Pattern_fail',
                  'PR_peeling', 'Seam', 'Reference', 'Surface_particle', 'Burried_particle', 'Cu_diffuse',
                  'Prelayer_defect_developed',
                  'Void', 'Residue', 'Scratch']

    categories_to_id = dict((c, i) for i, c in enumerate(categories))
    id_to_categories = dict((v, k) for k, v in categories_to_id.items())

    DEVICE = torch.device("cpu")

    # preprocess
    transform_test = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            # transforms.Normalize(mean=[0.42368844, 0.42368844, 0.42368844], std=[0.14975642, 0.14975642, 0.14975642])
            transforms.Normalize(mean=[0.40272543], std=[0.13901867])
        ])


    image_class = Image.open(class_image).convert("L")
    image_class = transform_test(image_class)
    image_class = torch.unsqueeze(image_class, dim=0).to(DEVICE).numpy()

    image_defect = Image.open(defect_image).convert("L")
    image_defect = transform_test(image_defect)
    image_defect = torch.unsqueeze(image_defect, dim=0).to(DEVICE).numpy()

    inputs_dict = {
        'class_images': image_class,
        'defect_images': image_defect
    }

    # Predict
    session = onnxruntime.InferenceSession(onnx_weight)
    # outputs = session.run(None, inputs_dict)
    outputs = None
    start_time = time.time()
    for i in range(100):
        outputs = session.run(None, inputs_dict)
        out = outputs[0]
        out = torch.from_numpy(out)
        _, pred = torch.max(out.data, 1)
    end_time = time.time()
    print((end_time - start_time) / 100)
    out = outputs[0]

    print(f"model output:{out}")
    out = torch.from_numpy(out)
    _, pred = torch.max(out.data, 1)
    print(f"defect id:{pred}")
    res = id_to_categories[pred.data.item()] # defect name
    print(f"defect name:{res}")
    # print('Image Name:{} \npredict:{}'.format(imagepath, id_to_categories[pred.data.item()]))
    return res

if __name__ == '__main__':
    class_image = 'origin_class.jpg'
    defect_image = 'origin_defect.jpg'
    onnx_weight = 'BiFNet_48_33125.onnx'
    predictClassDefectImages(class_image, defect_image, onnx_weight)

    '''
    model output:[[-12.145201  -17.541649   -7.728839  -14.660625  -17.136967  -14.414879
      -16.029175  -15.560219   -6.7670693 -12.183565   -9.317295   -2.4804323
      -15.176077  -10.80758     5.356476  -10.675158 ]]
    defect id:tensor([14])
    defect name:Residue
    '''
